!gdown 1T_RPkPP0jeWwK8L1UrmBw8V30eD7v6Ql

Downloading...
From (original): https://drive.google.com/uc?id=1T_RPkPP0jeWwK8L1UrmBw8V30eD7v6Ql
From (redirected): https://drive.google.com/uc?id=1T_RPkPP0jeWwK8L1UrmBw8V30eD7v6Ql&confirm=t&uuid=2085d32f-a224-4c84-b7a5-6dfb511a68d1
To: /content/PH2Dataset.rar
100% 162M/162M [00:01<00:00, 125MB/s]

get_ipython().system_raw("unrar x PH2Dataset.rar")

images = []
lesions = []
from skimage.io import imread
import os
root = 'PH2Dataset'

for root, dirs, files in os.walk(os.path.join(root, 'PH2 Dataset images')):
    if root.endswith('_Dermoscopic_Image'):
        images.append(imread(os.path.join(root, files[0])))
    if root.endswith('_lesion'):
        lesions.append(imread(os.path.join(root, files[0])))

from skimage.transform import resize
size = (256, 256)
X = [resize(x, size, mode='constant', anti_aliasing=True,) for x in images]
Y = [resize(y, size, mode='constant', anti_aliasing=False) > 0.5 for y in lesions]

import numpy as np
X = np.array(X, np.float32)
Y = np.array(Y, np.float32)
print(f'Loaded {len(X)} images')

Loaded 200 images

import matplotlib.pyplot as plt
from IPython.display import clear_output

plt.figure(figsize=(18, 6))
for i in range(6):
    plt.subplot(2, 6, i+1)
    plt.axis("off")
    plt.imshow(X[i])

    plt.subplot(2, 6, i+7)
    plt.axis("off")
    plt.imshow(Y[i])
plt.show();

ix = np.random.choice(len(X), len(X), False)
tr, val, ts = np.split(ix, [100, 150])

print(len(tr), len(val), len(ts))

100 50 50

from torch.utils.data import DataLoader
batch_size = 25
train_dataloader = DataLoader(list(zip(np.rollaxis(X[tr], 3, 1), Y[tr, np.newaxis])),
                     batch_size=batch_size, shuffle=True)
valid_dataloader = DataLoader(list(zip(np.rollaxis(X[val], 3, 1), Y[val, np.newaxis])),
                      batch_size=batch_size, shuffle=False)
test_dataloader = DataLoader(list(zip(np.rollaxis(X[ts], 3, 1), Y[ts, np.newaxis])),
                     batch_size=batch_size, shuffle=False)

loaders = {'train':train_dataloader, 'val': valid_dataloader}

import torch
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(DEVICE)

cuda

torch.manual_seed(42)
torch.cuda.manual_seed(42)
torch.cuda.manual_seed_all(42)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

!pip install -q torchmetrics

from torchmetrics import JaccardIndex

iou_score = JaccardIndex(threshold=0.5, task="binary", average='none').to(DEVICE)

import torch.nn.functional as F
import torch.nn as nn

bce_torch = nn.BCELoss(reduction='sum') # (sigmoid(y_pred), y_real)
bce_torch_with_logits = nn.BCEWithLogitsLoss(reduction='sum')

def bce_loss(y_pred, y_real):
  return torch.sum(y_pred - y_pred*y_real + torch.log(1 + torch.exp(-y_pred)))

def bce_true(y_pred, y_real):
  y = y_real
  p = 1 / (1 + torch.exp(-y_pred))
  return -torch.sum(y * torch.log(p) + (1 - y) * torch.log(1 - p))

y_pred = torch.randn(3, 2, requires_grad=False)
y_true = torch.rand(3, 2, requires_grad=False)

print(f'BCE loss from scratch bce_loss             = {bce_loss(y_pred, y_true)}')
print(f'BCE loss честно посчитанный                = {bce_true(y_pred, y_true)}')
print(f'BCE loss from torch bce_torch              = {bce_torch(torch.sigmoid(y_pred), y_true)}')
print(f'BCE loss from torch with logits bce_torch  = {bce_torch_with_logits(y_pred, y_true)}')

BCE loss from scratch bce_loss             = 4.616026401519775
BCE loss честно посчитанный                = 4.616025924682617
BCE loss from torch bce_torch              = 4.616025924682617
BCE loss from torch with logits bce_torch  = 4.616025924682617

assert np.isclose(bce_loss(y_pred, y_true), bce_torch(torch.sigmoid(y_pred), y_true))
assert np.isclose(bce_loss(y_pred, y_true), bce_torch_with_logits(y_pred, y_true))
assert np.isclose(bce_true(y_pred, y_true), bce_torch(torch.sigmoid(y_pred), y_true))
assert np.isclose(bce_true(y_pred, y_true), bce_torch_with_logits(y_pred, y_true))

y_pred = torch.randn((2, 1, 3, 3), requires_grad=False)
y_true = torch.randint(0, 2, (2, 1, 3, 3))

print(f'BCE loss from scratch bce_loss            = {bce_loss(y_pred, y_true)}')
print(f'BCE loss честно посчитанный               = {bce_true(y_pred, y_true)}')
print(f'BCE loss from torch bce_torch             = {bce_torch(torch.sigmoid(y_pred), y_true.to(torch.float))}')
print(f'BCE loss from torch with logits bce_torch = {bce_torch_with_logits(y_pred, y_true.to(torch.float))}')

BCE loss from scratch bce_loss            = 14.737800598144531
BCE loss честно посчитанный               = 14.737801551818848
BCE loss from torch bce_torch             = 14.737801551818848
BCE loss from torch with logits bce_torch = 14.737800598144531

assert np.isclose(bce_loss(y_pred, y_true), bce_torch(torch.sigmoid(y_pred), y_true.to(torch.float)))
assert np.isclose(bce_loss(y_pred, y_true), bce_torch_with_logits(y_pred, y_true.to(torch.float)))
assert np.isclose(bce_true(y_pred, y_true), bce_torch(torch.sigmoid(y_pred), y_true.to(torch.float)))
assert np.isclose(bce_true(y_pred, y_true), bce_torch_with_logits(y_pred, y_true.to(torch.float)))

!gdown --folder 1EX0RW1TRQVkLmR1h6miCQqyhYPFyg28M

Retrieving folder contents
Processing file 1--WxvBdpMn_NOmYPf3a4au8MHzfx5baC labels.pt
Processing file 1-0A7_CS_vKiSCkgIDJ4joThCEcFedA3I logits.pt
Retrieving folder contents completed
Building directory structure
Building directory structure completed
Downloading...
From: https://drive.google.com/uc?id=1--WxvBdpMn_NOmYPf3a4au8MHzfx5baC
To: /content/for_asserts/labels.pt
100% 1.18k/1.18k [00:00<00:00, 5.16MB/s]
Downloading...
From: https://drive.google.com/uc?id=1-0A7_CS_vKiSCkgIDJ4joThCEcFedA3I
To: /content/for_asserts/logits.pt
100% 1.18k/1.18k [00:00<00:00, 519kB/s]
Download completed

path_to_dummy_samples = '/content/for_asserts'
dummpy_sample = {'logits': torch.load(f'{path_to_dummy_samples}/logits.pt'),
                 'labels': torch.load(f'{path_to_dummy_samples}/labels.pt')}
dummpy_sample['labels'] = dummpy_sample['labels'].to(DEVICE)
dummpy_sample['logits'] = dummpy_sample['logits'].to(DEVICE)

fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize= (10,3*10))

ax1.imshow(dummpy_sample['labels'].squeeze(1)[0].cpu())
ax1.set_title("Original")

ax2.imshow(dummpy_sample['logits'].sigmoid().squeeze(1)[0].cpu())
for (j,i),label in np.ndenumerate(dummpy_sample['logits'].sigmoid().squeeze(1)[0].cpu()):
    if label < 0.5:
        color = 'white'
    else:
        color = 'black'
    ax2.text(i,j,round(label,3), color=color, ha='center',va='center')

ax2.set_title("Predicted Probabilities")

ax3.imshow((dummpy_sample['logits'].sigmoid() > 0.5).squeeze(1)[0].cpu())
ax3.set_title("Predicted Mask")
plt.show()

bce_loss_score = bce_loss(dummpy_sample['logits'].cpu(), dummpy_sample['labels'].cpu())
bce_true_score = bce_true(dummpy_sample['logits'].cpu(), dummpy_sample['labels'].cpu())
bce_torch_score = bce_torch(torch.sigmoid(dummpy_sample['logits'].cpu()), dummpy_sample['labels'].cpu().float())
bce_torch_with_logits_score = bce_torch_with_logits(dummpy_sample['logits'].cpu(), dummpy_sample['labels'].cpu().float())
assert np.isclose(bce_loss_score, bce_torch_score)
assert np.isclose(bce_loss_score, bce_torch_with_logits_score)
assert np.isclose(bce_true_score, bce_torch_score)
assert np.isclose(bce_true_score, bce_torch_with_logits_score)

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
import torch.optim as optim
from time import time

from matplotlib import rcParams
rcParams['figure.figsize'] = (15,4)

model_vgg16 = models.vgg16(weights = models.VGG16_Weights.IMAGENET1K_V1)

model_vgg16

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

# Параметрами блока будут:
# - количество каналов на входе
# - количество каналов на выходе
# - глубина блока (2 или 3, по количеству конволюционных слоев)
# - kernel_size и padding
#
class EncoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels, depth, kernel_size = 3, padding = 1):
        super(EncoderBlock, self).__init__() # инициируем экземляр класса, наследующего от nn.Module
        self.layers = nn.ModuleList() # в self.layers будем добавлять слои блока
        # дальше реализуем то, что на картинке выше обозначено Conv + Batch Normalization + ReLU
        self.layers.append(nn.Conv2d(in_channels = in_channels, out_channels = out_channels, kernel_size = kernel_size, padding = padding))
        self.layers.append(nn.BatchNorm2d(out_channels))
        self.layers.append(nn.ReLU(inplace=True))

        # цикл for помогает использовать один код для блоков как с глубиной 2, так и с глубиной 3
        for i in range(depth-1):
            self.layers.append(nn.Conv2d(in_channels = out_channels, out_channels = out_channels, kernel_size = kernel_size, padding = padding))
            self.layers.append(nn.BatchNorm2d(out_channels))
            self.layers.append(nn.ReLU(inplace=True))

        self.maxpooling = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True) #добавляем MaxPool с индексами для последующего Unpooling

    # Обратите внимание: на вход метод forward() получает карту признаков (х),
    # а возвращает карту признаков и индексы для последующего Unpooling
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        size = x.size()
        x, indices = self.maxpooling(x)
        return x, indices, size

class DecoderBlock(nn.Module):
  def __init__(self, in_channels, out_channels, depth, kernel_size = 3, padding = 1):
    super().__init__()

    self.layers = nn.ModuleList()
    for _ in range(depth):
      self.layers.append(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=padding))
      self.layers.append(nn.BatchNorm2d(out_channels))
      self.layers.append(nn.ReLU(inplace=True))
      in_channels = out_channels

    self.maxunpooling = nn.MaxUnpool2d(kernel_size=2, stride=2)

  def forward(self, x, indices, output_size):
    x = self.maxunpooling(x, indices, output_size)

    for layer in self.layers:
      x = layer(x)
    return x

class SegNet(nn.Module):
    def __init__(self, in_channels=3, out_channels = 1, num_features = 64) -> None:
        super(SegNet, self).__init__()

        # Encoder
        self.encoder0 = EncoderBlock(in_channels, num_features, depth=2)
        self.encoder1 = EncoderBlock(num_features, num_features * 2, depth=2)
        self.encoder2 = EncoderBlock(num_features * 2, num_features * 4, depth=3)
        self.encoder3 = EncoderBlock(num_features * 4, num_features * 8, depth=3)

        # Encoder bottleneck - количество каналов на входе и на выходе одинаково
        self.encoder4 = EncoderBlock(num_features * 8, num_features * 8, depth=3)

        # Decoder bottleneck
        self.decoder0 = DecoderBlock(num_features * 8, num_features * 8, depth=3)

        # Decoder
        self.decoder1 = DecoderBlock(num_features * 8, num_features * 4, depth=3)
        self.decoder2 = DecoderBlock(num_features * 4, num_features * 2, depth=3)
        self.decoder3 = DecoderBlock(num_features * 2, num_features, depth=2)
        self.decoder4 = DecoderBlock(num_features, num_features, depth=2)
        self.final = nn.Conv2d(num_features, out_channels, kernel_size=1)

    def forward(self, x):
        # encoder
        e0, ind0, size0 = self.encoder0(x)
        e1, ind1, size1 = self.encoder1(e0)
        e2, ind2, size2 = self.encoder2(e1)
        e3, ind3, size3 = self.encoder3(e2)
        e4, ind4, size4 = self.encoder4(e3)

        # Decoder
        d0 = self.decoder0(e4, ind4, size4)
        d1 = self.decoder1(d0, ind3, size3)
        d2 = self.decoder2(d1, ind2, size2)
        d3 = self.decoder3(d2, ind1, size1)
        d4 = self.decoder4(d3, ind0, size0)
        output = self.final(d4)
        return output  # no activation

from tqdm.notebook import tqdm

def fit_one_epoch(model, train_dataloader, optimizer, loss_func):
  '''
  args:
    model - модель для обучения
    train_dataloader - loader с выборкой для обучения модели
    optimizer - оптимизатор, взятый из модуля `torch.optim`
    loss_func - функция потерь, взятая из модуля `torch.nn`

  функция возвращает метрику accuracy по эпохе на данных из train_dataloader
  '''

  model.train()
  avg_loss = 0
  visualized = False

  for X_batch, y_batch in tqdm(train_dataloader):
    X_batch = X_batch.to(DEVICE)
    y_batch = y_batch.to(DEVICE)

    optimizer.zero_grad()
    outp = model(X_batch)

    prob = torch.sigmoid(outp)
    y_pred = (prob > 0.5).long()
    if not visualized:
      visualized = True
      visualize(X_batch, y_batch, y_pred)

    loss = loss_func(outp, y_batch)
    loss.backward()
    optimizer.step()

    avg_loss += loss.item()

  return avg_loss / len(train_dataloader)

def eval_one_epoch(model, val_dataloader, loss_func):
  '''
  args:
    model - модель для обучения
    val_dataloader - loader с валидационной/тестовой выборкой
  '''
  iou_score = JaccardIndex(threshold=0.5, task="binary", average='none').to(DEVICE)

  model.eval()
  avg_loss = 0
  avg_iou = 0

  visualized = False

  with torch.no_grad():
    for X_batch, y_batch in tqdm(val_dataloader):
      X_batch = X_batch.to(DEVICE)
      y_batch = y_batch.to(DEVICE)

      outp = model(X_batch)
      prob = torch.sigmoid(outp)
      y_pred = (prob > 0.5).long()

      loss = loss_func(outp, y_batch)
      iou = iou_score(y_pred, y_batch)

      avg_loss += loss.item()
      avg_iou += iou

      if not visualized:
        visualized = True
        visualize(X_batch, y_batch, y_pred)

    avg_loss = avg_loss / len(val_dataloader)
    avg_iou = avg_iou / len(val_dataloader)

  return avg_loss, avg_iou

def visualize(X_batch, y_batch, pred, n=1):
  batch_size = X_batch.shape[0]

  plt.figure(figsize=(10, 3*n))
  for i in range(n):
    img = X_batch[i].permute(1,2,0).cpu().numpy()
    true_mask = y_batch[i].cpu().squeeze().numpy()
    pred_mask = pred[i].squeeze().cpu().numpy()

    plt.subplot(n, 3, i*3 + 1)
    plt.title("Image")
    plt.imshow(img)
    plt.axis("off")

    plt.subplot(n, 3, i*3 + 2)
    plt.title("GT mask")
    plt.imshow(true_mask, cmap="nipy_spectral")
    plt.axis("off")

    plt.subplot(n, 3, i*3 + 3)
    plt.title("Pred mask")
    plt.imshow(pred_mask, cmap="nipy_spectral")
    plt.axis("off")

  plt.show()

def train_func(model, num_epochs, dataloaders, optimizer, loss_func):
  '''
  args:
    model - модель для обучения
    num_epochs - количество эпох
    dataloaders - словарь loader'ов с обучающей и валидационной выборками
    optimizer - оптимизатор, взятый из модуля `torch.optim`
    loss_func - функция потерь, взятая из модуля `torch.nn`

  функция возвращает loss на обучающей и валидационной выборках на каждой эпохе, а также метрику IoU на валидационной выборке
  '''
  model = model.to(DEVICE)
  score = {"train_loss": [], "val_loss": [], "val_iou": []}
  for epoch in range(num_epochs):
    print(f"\nEpoch: {epoch+1}")

    loss_train = fit_one_epoch(model = model, train_dataloader = dataloaders['train'], optimizer = optimizer, loss_func = loss_func)
    print(f"Loss train: {loss_train}\n")

    loss_val, iou_val = eval_one_epoch(model = model, val_dataloader = dataloaders['val'], loss_func=loss_func)
    print(f"Loss valid: {loss_val}\n")
    print(f"IoU valid: {iou_val}\n")


    score['train_loss'].append(loss_train)
    score['val_loss'].append(loss_val)
    score['val_iou'].append(iou_val)
  return model, score

model_baseline = SegNet()

criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model_baseline.parameters(), lr = 1e-3)

model_baseline, score_baseline = train_func(model_baseline, 50, loaders, optimizer, criterion)

Epoch: 1

  0%|          | 0/4 [00:00<?, ?it/s]

Loss train: 0.7088902145624161

  0%|          | 0/2 [00:00<?, ?it/s]

Loss valid: 0.6806586980819702

IoU valid: 0.0


Epoch: 2

  0%|          | 0/4 [00:00<?, ?it/s]

Loss train: 0.6382251232862473

  0%|          | 0/2 [00:00<?, ?it/s]

Loss valid: 0.6685355007648468

IoU valid: 0.0


Epoch: 3

  0%|          | 0/4 [00:00<?, ?it/s]

def test(model, test_dataloader):
  model.eval()

  iou_score = JaccardIndex(threshold=0.5, task="binary", average='none').to(DEVICE)
  avg_iou = 0
  with torch.no_grad():
    for X_batch, Y_batch in test_dataloader:
      X_batch = X_batch.to(DEVICE)
      Y_batch = Y_batch.to(DEVICE)

      outp = model(X_batch)
      prob = torch.sigmoid(outp)
      y_pred = (prob > 0.5).long()
      avg_iou += iou_score(y_pred, Y_batch)

  return avg_iou / len(test_dataloader)

test_score = test(model_baseline, test_dataloader)

test_score.item()

0.8248778581619263

def dice_score(logits: torch.Tensor, labels: torch.Tensor, threshold: float = 0.5):
  '''
  Это именно метрика, не лосс.
  '''
  eps = 1e-8

  prob = torch.sigmoid(logits)
  preds = (prob > threshold).int()

  TP = (preds * labels).sum()
  FP = (preds * (1 - labels)).sum()
  FN = ((1 - preds) * labels).sum()

  score = (2 * TP + eps) / (2 * TP + FP + FN + eps)

  return score

from torchmetrics.segmentation import DiceScore

dice = DiceScore(num_classes=1, average='micro').to(DEVICE)
dice(dummpy_sample['logits'].sigmoid() > 0.5, dummpy_sample['labels'].int())

/usr/local/lib/python3.12/dist-packages/torchmetrics/utilities/prints.py:43: UserWarning: DiceScore metric currently defaults to `average=micro`, but will change to`average=macro` in the v1.9 release. If you've explicitly set this parameter, you can ignore this warning.
  warnings.warn(*args, **kwargs)

tensor(0.6667, device='cuda:0')

assert dice(dummpy_sample['logits'].sigmoid()>0.5, dummpy_sample['labels'].to(int)) == dice_score(dummpy_sample['logits'], dummpy_sample['labels'])

def dice_loss(logits: torch.Tensor, labels: torch.Tensor):
  eps = 1e-8

  probs = torch.sigmoid(logits)

  TP = (probs * labels).sum()
  FP = (probs * (1 - labels)).sum()
  FN = ((1 - probs) * labels).sum()

  dice = (2 * TP + eps) / (2 * TP + FP + FN + eps)
  return 1 - dice.mean()

# проверьте, что у вас установлена библиотека
!pip install -q segmentation-models-pytorch

from segmentation_models_pytorch.losses import DiceLoss
dice_loss_torch = DiceLoss(mode='binary')
dice_loss_torch(dummpy_sample['logits'], dummpy_sample['labels'])

tensor(0.5756, device='cuda:0')

dice_loss(dummpy_sample['logits'], dummpy_sample['labels'])

tensor(0.5756, device='cuda:0')

assert dice_loss_torch(dummpy_sample['logits'], dummpy_sample['labels'].to(int)) == dice_loss(dummpy_sample['logits'], dummpy_sample['labels'])

def focal_loss(y_real, y_pred, eps = 1e-6, gamma = 2):
  p = torch.sigmoid(y_pred) * y_real + (1 - torch.sigmoid(y_pred))*(1 - y_real)
  bce_loss = -(y_real * torch.log(torch.sigmoid(y_pred) + eps) + (1 - y_real) * torch.log(1 - torch.sigmoid(y_pred) + eps))
  loss = (1 - p)**gamma * bce_loss
  return loss.sum()

from torchvision.ops import sigmoid_focal_loss
sigmoid_focal_loss(dummpy_sample['logits'], dummpy_sample['labels'], alpha=-1, gamma=2, reduction='sum').item()

3.616123676300049

assert torch.allclose(sigmoid_focal_loss(dummpy_sample['logits'], dummpy_sample['labels'], alpha=-1, gamma=2, reduction='sum'),
                      focal_loss(dummpy_sample['labels'], dummpy_sample['logits'], gamma=2.0),
                      rtol=1e-5,
                      atol=1e-8)

from typing import Tuple
import torch
import torch.nn.functional as F

def make_gaussian_kernel(k: int, sigma: float):
  assert k % 2 == 1, "kernel size must be odd"
  half = k // 2
  xs = torch.arange(-half, half+1, dtype=torch.float32)
  ys = xs.view(-1, 1)
  kernel = torch.exp(-(xs**2 + ys**2) / (2 * sigma**2))
  kernel = kernel / kernel.sum()
  return kernel

class BinaryStructuralSimilarityLoss(nn.Module):
  def __init__(self, window_size=3, tau=0.1, lambda_ce=0.9, c=0.01, eps=1e-8, gauss_sigma=1.5):
    super().__init__()
    self.k = window_size
    self.tau = tau
    self.lambda_ce = lambda_ce
    self.c = c
    self.eps = eps
    self.gauss_sigma = gauss_sigma

  def __get_kernel(self):
    return make_gaussian_kernel(self.k, self.gauss_sigma)

  def count_abs_structural_error(self, probs: torch.Tensor, labels: torch.Tensor):
    B = probs.shape[0]
    p_unf, y_unf = self.make_patches(probs, labels) # (B, k*k, H*W)

    kernel = self.__get_kernel()

    statistics_p = self.compute_statistics_over_patch(probs, kernel) # (B,1,H,W)
    statistics_y = self.compute_statistics_over_patch(labels, kernel)

    mean_p, sigma_p = self.adjust_dims(statistics_p, B) # (B, k*k, H*W)
    mean_y, sigma_y = self.adjust_dims(statistics_y, B)

    z_p = (p_unf - mean_p + self.c) / (sigma_p + self.c)
    z_y = (y_unf - mean_y + self.c) / (sigma_y + self.c)

    error = torch.abs(z_y - z_p).sum(dim=1)  # (B, H*W)
    return error

  def make_patches(self, probs: torch.Tensor, labels: torch.Tensor):
    pad = self.k // 2
    p_unf = F.unfold(probs, kernel_size=self.k, padding=pad) # (B, k*k, H*W)
    l_unf = F.unfold(labels, kernel_size=self.k, padding=pad) # (B, k*k, H*W)

    return p_unf, l_unf

  def compute_statistics_over_patch(self, labels: torch.Tensor, kernel: torch.Tensor, eps: float = 1e-8):
    pad = self.k // 2
    kernel = kernel.view(1, 1, self.k, self.k).to(labels.device)

    mean = F.conv2d(labels, kernel, padding=pad)  # (B,1,H,W)
    E2 = F.conv2d(labels * labels, kernel, padding=pad)
    var = torch.clamp(E2 - mean * mean, min=0.0)
    sigma = torch.sqrt(var + eps)  # (B,1,H,W)

    return mean, sigma

  def adjust_dims(self, statistics: Tuple[torch.Tensor, torch.Tensor], num_batches: int):
    k = self.k
    B = num_batches
    mean_exp = statistics[0].view(B, 1, -1)  # (B, 1, H*W)
    sigma_exp = statistics[1].view(B, 1, -1)  # (B, 1, H*W)
    mean_exp = mean_exp.repeat(1, k*k, 1)  # (B, k*k, H*W)
    sigma_exp = sigma_exp.repeat(1, k*k, 1)
    return mean_exp, sigma_exp

  def forward(self, logits: torch.Tensor, labels: torch.Tensor):
    probs = torch.sigmoid(logits)
    error = self.count_abs_structural_error(probs, labels)
    indicator = (error > self.tau * error.max(dim=1, keepdim=True)[0]).float()

    bce_full = F.binary_cross_entropy_with_logits(logits, labels, reduction='mean')

    bce_pixel = F.binary_cross_entropy_with_logits(logits, labels, reduction='none').view(probs.shape[0], -1) # (B, H*W)
    L_ssl_pixel = error * indicator * bce_pixel # (B, H*W)

    M = indicator.sum(dim=1).clamp(min=1) # (B,)
    L_ssl_batch = (L_ssl_pixel.sum(dim=1) / M) # (B,)
    L_ssl = L_ssl_batch.mean()

    L_all = self.lambda_ce * bce_full + (1 - self.lambda_ce) * L_ssl
    return L_all

ssloss = BinaryStructuralSimilarityLoss()
ssloss(dummpy_sample['logits'], dummpy_sample['labels'])

tensor(0.5462, device='cuda:0')

logits = torch.tensor([[[[
    -1.,  1.,  2.,  1., -1.,
     1., -2., -3., -2.,  1.,
     2., -3., -4., -3.,  2.,
     1., -2., -3., -2.,  1.,
    -1.,  1.,  2.,  1., -1.,
]]]], dtype=torch.float32)

labels = torch.tensor([[[[
    0., 1., 1., 1., 0.,
    1., 0., 0., 0., 1.,
    1., 0., 0., 0., 1.,
    1., 0., 0., 0., 1.,
    0., 1., 1., 1., 0.,
]]]], dtype=torch.float32)

print("SS Loss:", ssloss(logits, labels))
print("Dice Loss:", dice_loss(logits, labels))
print("Focal Loss:", focal_loss(logits, labels))
print("BCE Loss:", bce_loss(logits, labels))

SS Loss: tensor(0.2240)
Dice Loss: tensor(0.1897)
Focal Loss: tensor(2.3315)
BCE Loss: tensor(4.9871)

labels = torch.tensor([[
    [0,0,0,0,0,0,0,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,0,0,0,0,0,0,0],
    [0,0,0,0,0,0,0,0,0,0]
]], dtype=torch.float)

logits = torch.where(labels==1, torch.tensor(10.0), torch.tensor(-10.0)).unsqueeze(0)
# shape: (1,1,10,10)
labels = labels.unsqueeze(0)

print("SS Loss:", ssloss(logits, labels))
print("Dice Loss:", dice_loss(logits, labels))
print("Focal Loss:", focal_loss(logits, labels))
print("BCE Loss:", bce_loss(logits, labels))

SS Loss: tensor(4.1325e-05)
Dice Loss: tensor(8.1062e-05)
Focal Loss: tensor(-3668.2905)
BCE Loss: tensor(0.0046)

labels = torch.tensor([[
    [0,0,0,0,0,0,0,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,1,1,1,1,0,0,0],
    [0,0,0,0,0,0,0,0,0,0],
    [0,0,0,0,0,0,0,0,0,0]
]], dtype=torch.float)

logits = torch.where(labels==1, torch.tensor(-10.0), torch.tensor(+10.0)).unsqueeze(0)
labels = labels.unsqueeze(0)

print("SS Loss:", ssloss(logits, labels))
print("Dice Loss:", dice_loss(logits, labels))
print("Focal Loss:", focal_loss(logits, labels))
print("BCE Loss:", bce_loss(logits, labels))

SS Loss: tensor(22.0742)
Dice Loss: tensor(1.0000)
Focal Loss: tensor(9086.8047)
BCE Loss: tensor(1000.0046)

model_dice = SegNet()
criterion = DiceLoss(mode='binary')
optimizer = torch.optim.Adam(model_dice.parameters(), lr = 1e-3)
model_dice, score_dice = train_func(model_dice, 50, loaders, optimizer, criterion)

def focal_loss_wrapper(logits, labels):
    loss = sigmoid_focal_loss(logits, labels, reduction='none')
    return loss.mean()

model_focal = SegNet()
criterion = focal_loss_wrapper
optimizer = torch.optim.Adam(model_focal.parameters(), lr = 1e-3)
model_focal, score_focal = train_func(model_focal, 50, loaders, optimizer, criterion)

def ss_loss_wrapper(logits, labels):
  loss = ssloss(logits, labels)
  return loss.mean()

model_ssl = SegNet()
criterion = ss_loss_wrapper
optimizer = torch.optim.Adam(model_ssl.parameters(), lr = 1e-3)
model_ssl, score_ssl = train_func(model_ssl, 50, loaders, optimizer, criterion)

def to_float_list(x):
    return [float(v) for v in x]  # работает и для CUDA tensors

score_baseline['val_iou'] = to_float_list(score_baseline['val_iou'])
score_dice['val_iou']      = to_float_list(score_dice['val_iou'])
score_focal['val_iou']     = to_float_list(score_focal['val_iou'])
score_ssl['val_iou']        = to_float_list(score_ssl['val_iou'])

plt.plot(range(1, 51), score_baseline['val_iou'], label='score baseline')
plt.plot(range(1, 51), score_dice['val_iou'], label='score dice')
plt.plot(range(1, 51), score_focal['val_iou'], label='score focal')
plt.plot(range(1, 51), score_ssl['val_iou'], label='score ssl')


plt.title('IoU на валидационной выборке')
plt.xlabel('epoch')
plt.ylabel('IoU score')
plt.legend()

<matplotlib.legend.Legend at 0x78bc651fb410>

plt.plot(range(1, 51), score_baseline['val_loss'], label='score baseline')
plt.plot(range(1, 51), score_dice['val_loss'], label='score dice')
plt.plot(range(1, 51), score_focal['val_loss'], label='score focal')
plt.plot(range(1, 51), score_ssl['val_loss'], label='score ssl')

plt.ylim(0, 5)

plt.title('Loss на валидационной выборке')
plt.xlabel('epoch')
plt.ylabel('Loss')
plt.legend()

<matplotlib.legend.Legend at 0x78bc64063200>

test_dice = test(model_dice, test_dataloader)
test_focal = test(model_focal, test_dataloader)
test_ssl = test(model_ssl, test_dataloader)

print("Модель с Dice Loss имеет Iou:", test_dice.cpu().item())
print("Модель с Focal Loss имеет Iou:", test_focal.cpu().item())
print("Модель с SS Loss имеет Iou:", test_ssl.cpu().item())

Модель с Dice Loss имеет Iou: 0.8408982753753662
Модель с Focal Loss имеет Iou: 0.8346421718597412
Модель с SS Loss имеет Iou: 0.681615948677063

import torch.nn.functional as F
import torch.nn as nn

class UNetEncoder(nn.Module):
  def __init__(self, in_channels, out_channels):
    super().__init__()

    self.conv = nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
    )
    self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

  def forward(self, x):
    x_conv = self.conv(x)
    x_down = self.pool(x_conv)
    return x_conv, x_down

class UNetDecoder(nn.Module):
  def __init__(self, in_channels, out_channels, padding=1):
    super().__init__()

    self.up_conv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)

    self.conv = nn.Sequential(
        nn.Conv2d(out_channels * 2, out_channels, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
    )

  def forward(self, x, skip):
    x = self.up_conv(x)
    x = torch.cat([skip, x], dim=1)
    return self.conv(x)

class UNet(nn.Module):
    def __init__(self, n_class=1):
        super().__init__()

        self.e0 = UNetEncoder(3, 64)
        self.e1 = UNetEncoder(64, 128)
        self.e2 = UNetEncoder(128, 256)
        self.e3 = UNetEncoder(256, 512)
        self.e4 = UNetEncoder(512, 1024)

        self.d0 = UNetDecoder(1024, 512)
        self.d1 = UNetDecoder(512, 256)
        self.d2 = UNetDecoder(256, 128)
        self.d3 = UNetDecoder(128, 64)

        self.final_conv = nn.Conv2d(64, n_class, kernel_size=1)

    def forward(self, x):
        x0_conv, x0_down = self.e0(x)
        x1_conv, x1_down = self.e1(x0_down)
        x2_conv, x2_down = self.e2(x1_down)
        x3_conv, x3_down = self.e3(x2_down)
        x4_conv, x4_down = self.e4(x3_down)

        x = self.d0(x4_conv, x3_conv)
        x = self.d1(x, x2_conv)
        x = self.d2(x, x1_conv)
        x = self.d3(x, x0_conv)

        output = self.final_conv(x)
        return output

unet_model = UNet().to(DEVICE)
criterion = focal_loss_wrapper
optimizer = torch.optim.Adam(unet_model.parameters(), lr = 1e-3)
unet_model, unet_score_focal = train_func(unet_model, 50, loaders, optimizer, criterion)

def dice_loss_wrapper(logits, labels):
    loss = dice_loss(logits, labels)
    return loss.mean()

unet_model_dice = UNet().to(DEVICE)
criterion = dice_loss_wrapper
optimizer = torch.optim.Adam(unet_model_dice.parameters(), lr = 3e-4)
unet_model_dice, unet_score_dice = train_func(unet_model_dice, 50, loaders, optimizer, criterion)

unet_model_bce = UNet().to(DEVICE)
criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(unet_model_bce.parameters(), lr = 1e-3)
unet_model_bce, unet_score_bce = train_func(unet_model_bce, 50, loaders, optimizer, criterion)

test_unet_dice = test(unet_model_dice, test_dataloader)
test_unet_focal = test(unet_model, test_dataloader)
test_unet_bce = test(unet_model_bce, test_dataloader)

print("Модель с Dice Loss имеет Iou:", test_unet_dice.cpu().item())
print("Модель с Focal Loss имеет Iou:", test_unet_focal.cpu().item())
print("Модель с BCE Loss имеет Iou:", test_unet_bce.cpu().item())

Модель с Dice Loss имеет Iou: 0.3579278588294983
Модель с Focal Loss имеет Iou: 0.025996873155236244
Модель с BCE Loss имеет Iou: 0.5349061489105225

def to_float_list(x):
    return [float(v) for v in x]

unet_score_bce['val_iou'] = to_float_list(unet_score_bce['val_iou'])
unet_score_dice['val_iou']      = to_float_list(unet_score_dice['val_iou'])
unet_score_focal['val_iou']     = to_float_list(unet_score_focal['val_iou'])

plt.plot(range(1, 51), unet_score_bce['val_iou'], label='score baseline')
plt.plot(range(1, 51), unet_score_dice['val_iou'], label='score dice')
plt.plot(range(1, 51), unet_score_focal['val_iou'], label='score focal')

plt.title('IoU на валидационной выборке для модели U-Net')
plt.xlabel('epoch')
plt.ylabel('IoU score')
plt.legend()

<matplotlib.legend.Legend at 0x7f73d739ae10>

Школа глубокого обучения ФПМИ МФТИ

Домашнее задание. Сегментация изображений

Шаг 1. Загрузка и подготовка данных¶

PyTorch DataLoader¶

Шаг 2. Метрика качества модели¶

IoU (intersection over union)¶

Задания: Построй свой первый бейзлайн!¶

Шаг 3. Loss функция - BCE [2 балла]¶

Вывод численно стабильной формулы BCE лосса [1 балл]¶

Реализуйте в коде оба варианта лосса [1 балл]¶

Шаг 4. Модель SegNet [2 балла]¶

Шаг 5. Тренировка модели [1 балл]¶

Шаг 6. Инференс [1 балл]¶

Задания: Мир других лоссов!¶

Пробуем другие функции потерь [2 балла]¶

Dice Loss¶

Focal Loss¶

[BONUS] Мир сегментационных лоссов [5 баллов]¶

Обучите SegNet на новых лоссах¶

Задание: Новая модель!¶

Модель U-Net [2 балла]¶

Обучите UNet¶